package com.itheima.datastructure.avltree;

/*
AVL树
1.二叉树在插入和删除的时候才有可能失去平衡
2.如果一个树的左右孩子高度差大于1，则此节点失衡，需要旋转
 */
public class AVLTree {
    static class AVLNode{
        int key;
        Object value;
        AVLNode left;
        AVLNode right;
        int height=1;//高度

        public AVLNode(int key,Object value) {
            this.key = key;
            this.value=value;
        }

        public AVLNode(int key) {
            this.key = key;
        }

        public AVLNode(int key, Object value, AVLNode left, AVLNode right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }
    }

    //求节点的高度值
    private int height(AVLNode node){
        return node==null?0:node.height;
    }

    //更新节点高度：就是左右子树节点高度最大值+1；
    private void updateHeight(AVLNode node){
        node.height=Integer.max(height(node.left),height(node.right))+1;
    }

    //平衡因子：（balance factor)=左子树高度-右子树高度
    /*
    bf==0 ，i，-1表示平衡
    bf>1 || bf<-1 表示不平衡
     */
    private int bf(AVLNode node){
        return height(node.left)-height(node.right);
    }

    /*
    四种失衡情况：
    LL:失衡节点左子树比右子树高，左子树的左子树比右子树高
    LR:失衡节点左子树比右子树高，左子树的右子树比左子树高
    RL:失衡节点右子树比左子树高，左子树的左子树比右子树高
    RR:失衡节点右子树比右子树高，左子树的右子树比左子树高
     */

    //右旋 参数：要旋转的节点，返回新的根节点
    private AVLNode rightRotate(AVLNode red){
        AVLNode yellow=red.left;
        AVLNode green=yellow.right;
        yellow.right=red;//上位
        red.left=green;//换爹
        updateHeight(red);
        updateHeight(yellow);
        return yellow;
    }

    //左旋  参数：要旋转的节点，返回新的根节点
    private AVLNode leftRotate(AVLNode red){
        AVLNode yellow=red.right;
        AVLNode green=yellow.left;
        yellow.left=red;
        yellow.right=green;
        updateHeight(red);
        updateHeight(yellow);
        return yellow;
    }

    //先左旋左子树，再右旋右子树
    private AVLNode leftRightRotate(AVLNode node){
        node.left=leftRotate(node.left);
        return rightRotate(node);
    }

    //先右旋右子树，再左旋左子树
    private AVLNode rightLeftRotate(AVLNode node){
        node.right=rightRotate(node.right);
        return leftRotate(node);
    }

    //检查节点是否失衡，重新平衡
    private  AVLNode balance(AVLNode node){
        if(node==null){
            return null;
        }
        int bf=bf(node);
        if(bf>1 && bf((node.left))>=0){ //LL
            return rightRotate(node);
        }else if(bf>1 && bf(node.left)<0){ //LR
            return leftRightRotate(node);
        }else if(bf<-1 && bf(node.right)>0){ //RL
            return rightLeftRotate(node);
        }else if(bf<-1 && bf(node.right)<=0){ //RR
            return leftRotate(node);
        }
        return node;
    }

    AVLNode root;
    //添加新节点操作
    public void put(int key,Object value){
        root=doPut(root,key,value);
    }
    private AVLNode doPut(AVLNode node,int key,Object value){
        //1.找到空位就创建新节点
        if(node==null){
            return new AVLNode(key,value);
        }
        //2.如果key已经存在就更新value
        if(key==node.key){
            node.value=value;
            return node;
        }
        //3.继续寻找
        if(key<node.key){
            node.left=doPut(node.left,key,value);
        }else{
            node.right=doPut(node.right,key,value);
        }
        //4.更新高度值
        updateHeight(node);
        //5.重新建立平衡关系
        return balance(node);
    }

    //删除操作
    public void remove(int key){
        root=doRemove(root,key);
    }
    private AVLNode doRemove(AVLNode node,int key){
        //1.node==null
        if(node==null){
            return null;
        }
        //2.没有找到key
        if(key<node.key){
            node.left=doRemove(node.left,key);
        }else if(key>node.key){
            node.right=doRemove(node.right,key);
        }else{
            //3.找到了key
            /*
            1.没有孩子
            2.只有一个孩子
            3.有两个孩子
             */
            if(node.left==null && node.right==null){
                return null;
            }else if(node.left==null){
                node=node.right;
            }else if(node.right==null){
                node=node.left;
            }else{
                AVLNode s=node.right;
                //寻找后继节点
                while(s.left!=null){
                    s=s.left;
                }
                s.right=doRemove(node.right,s.key);
                s.left=node.left;
                node=s;
            }
        }
        //4.更新高度值
        updateHeight(node);
        return balance(node);
    }

}
